`lmbda` selects between training on the dataset's own completions
(`lmbda=0.0`) and on completions the student generates (`lmbda=1.0`). The
trainer is scheduled to become always-on-policy, so pin both modes now: the
later removal of the off-policy path should be a deliberate deletion with a
visibly failing test, not a silent behaviour change.

Asserts the loss is recorded and that every parameter actually moved, using
`torch.equal` rather than a tolerance so a stalled update cannot pass. Uses a
higher learning rate than the default because gradients are tiny on the test
model and the default rate can stall the update, which would make the
assertion vacuous.

The divergence loss is reduced as `sum / num_items_in_batch`, so that count is
the loss denominator. It is currently wrong.

`transformers.Trainer` computes it from the raw dataloader batches, before
`_prepare_inputs` runs. Two things follow. `_RepeatBatchDataLoader` yields the
same generation batch once per accumulation step, so the count is
`gradient_accumulation_steps` times too large. And it is derived from the
dataset labels, even on steps whose completions are replaced by generated ones.

The effect is a silently down-scaled loss and gradient. Add an xfail test that
records what the trainer is actually handed during `train()` and compares it
with the tokens the loss is summed over. The existing coverage passes
`num_items_in_batch` explicitly, so it exercises the reduction rather than the
value, and cannot catch this.

Un-xfail when the generation buffer moves to the GRPO-style `_prepare_inputs`
and the denominator is derived from the completions actually trained on.

IW-OPD (Importance-Weighted On-Policy Distillation, 2606.22600) is a
teacher-guided policy-gradient method: sampled-token advantage
`A_t = sg[log π_T − log π_0]`, a prefix-decayed importance weight, and a
clipped-surrogate PPO inner loop. It fits neither the stable full-vocabulary
`DistillationTrainer` (supervised KD, no advantages) nor `GRPOTrainer` (which
has no teacher). So rather than delete it, move it to its own home.

- Fork the pre-refactor trainer/config into `trl/experimental/iw_opd/` as
  `IWOPDTrainer` / `IWOPDConfig` — a self-contained snapshot that keeps IW-OPD
  (and, incidentally, the other pre-refactor features) working. Its docstring
  states plainly that it is a frozen, unmaintained snapshot.
- Remove IW-OPD from the base `DistillationTrainer` (loss, config fields,
  rollout-logprob plumbing, the `logprobs=0` vLLM request, and the tests).
- Repoint the IW-OPD paper-index entry to `experimental.iw_opd` (the GKD and
  general on-policy-distillation entries stay on `DistillationTrainer`).
- Register `IWOPDTrainer` in the telemetry allowlist; the moved tests live in
  `tests/experimental/test_iw_opd_trainer.py`.

This merges what were two PRs (remove IW-OPD + drop its paper entry) into one
"move" so nothing is lost. `compute_loss` in the base still drops from 7
dispatch paths to 5.

The server-backed path (fetch per-token teacher logprobs from a vLLM server
instead of a local teacher forward) is a distinct use case that only ever runs
over a sparse top-k support. Move it out of the base trainer into a dedicated
experimental subclass so the base can converge on the local, full-vocabulary
objective.

New package `trl/experimental/server_distillation/`:
- `ServerDistillationTrainer(DistillationTrainer)` overriding `compute_loss`
  with the server path, plus the moved `_get_teacher_token_logprobs_from_server`,
  `_compute_server_sparse_top_1_divergence_loss`, `_compute_server_forward_kl_loss`
  and the module-level `build_teacher_request_inputs`. The shared sparse helpers
  (`_compute_sparse_top_1_divergence_loss`, `_get_reverse_kl_top_1_tokens`,
  `_add_tail_bucket`, `_jsd_divergence`, `_reduce_divergence_loss`) are inherited
  from the base, not copied — they are still used by the base local top-1 path.
- `ServerDistillationConfig(DistillationConfig)` adding `teacher_model_server_url`
  and carrying the server-only validations, now unconditional.

Base trainer:
- `__init__` drops the server branch; a `None` teacher just leaves
  `self.teacher_model = None`.
- `compute_loss` drops the server dispatch; the local path runs unconditionally.
- `_get_teacher_logits` drops the server NotImplementedError branch.
- config drops `use_teacher_server` / `teacher_model_server_url` and their checks.

Server tests move to `tests/experimental/test_server_distillation_trainer.py`.

The server-side infrastructure (`/get_sequence_logprobs/`, VLLMClient) is
untouched and still used by SDFT/SDPO.

`compute_loss` in the base goes from 5 dispatch paths to 2 (local top-1, local
full-vocab) plus the Liger path.

With the server path extracted, the base trainer no longer needs the mixed
top-1 support path: the local teacher always has full logits, so it can compute
the exact generalized JSD/KL directly.

- `compute_loss` local branch now always calls `generalized_jsd_loss` (which
  still honours `loss_top_k` for the top-k approximation); the
  `beta > 0 and loss_top_k == 1` special case and its `completion_tokens` slice
  are gone.
- `_compute_local_sparse_top_1_divergence_loss` is deleted.
- `_compute_sparse_top_1_divergence_loss` and `_get_reverse_kl_top_1_tokens`
  become server-only, so they move to `ServerDistillationTrainer`. The subclass
  now defines them directly rather than inheriting them.
- `reverse_kl_top_1_mode` only drove those helpers, so it moves from
  `DistillationConfig` to `ServerDistillationConfig` (field + the
  "must be one of" validation), and the `self.reverse_kl_top_1_mode` assignment
  moves from the base `__init__` to the subclass `__init__`. Its stale mentions
  in the `loss_top_k` docs are trimmed.
- The corresponding config test moves to the server test file.

Verified on CPU that the server sparse loss is numerically unchanged by the
relocation.

The base trainer always has full teacher logits, so it never needs the top-k
approximation — only the extracted server path (which sees just the teacher's
top-k logprobs) does.

- `generalized_jsd_loss` drops the `top_k` / `add_tail` parameters and the whole
  top-k support-selection block; it is now the pure full-vocabulary path.
- config `loss_top_k` / `loss_add_tail` move from `DistillationConfig` to
  `ServerDistillationConfig`; the `self.loss_top_k` / `self.loss_add_tail`
  assignments move from the base `__init__` to the subclass `__init__`; the
  base `compute_loss` call drops the two kwargs.
- `_add_tail_bucket` was only reachable from the top-k block and the server
  methods, so it moves to the server trainer module. `_jsd_divergence` stays
  shared in the base module (the base uses its dense branch, the server its
  masked branch).
- `test_loss_normalizes_by_num_items_in_batch` drops its `loss_top_k`
  parametrization (base is full-vocab only); the `_add_tail_bucket` import in
  the server test now comes from the server module.

Verified on CPU that `generalized_jsd_loss` still matches a naive full-vocab
reference for beta in {0, 0.5, 1} with and without labels, and that the server
sparse loss is numerically unchanged.

The base local-teacher loss compares full next-token distributions, so the only
real requirement is a shared vocabulary. Replace the tokenizer-loading /
`get_vocab()` comparison + warn + `resize_token_embeddings` dance with a strict
`config.vocab_size` equality check that raises (pointing at GOLD for the
cross-tokenizer case), mirroring GKD.

- `__init__` teacher setup no longer loads the teacher tokenizer, no longer
  requires `config._name_or_path`, and no longer resizes the teacher embeddings.
  The check is gated on `teacher_model is not None`, so subclasses that supply
  the teacher another way (ServerDistillationTrainer) are unaffected.
- Remove `_local_teacher_tokenizers_match`, `_raise_if_local_teacher_tokenizer_mismatch`,
  the `self._local_teacher_tokenizer_matches_student` flag, and the guard call at
  the top of `compute_loss`. Drop the now-unused `warnings` import.
- Add tests for the vocab-size mismatch (Qwen2 student + Llama teacher) and for
  passing `teacher_model_init_kwargs` alongside an already-instantiated teacher.

Note: `teacher_model_name_or_path` remains an unused config field for now; it is
removed later when the teacher becomes a constructor-only argument.

The distillation trainer is converging on always-on-policy training (the
GRPO-shaped stable target generates every batch). `lmbda`, which mixes
on-policy (student-generated) and off-policy (dataset) completions per
gradient-accumulation slice, is being removed.

Warn with FutureWarning when `lmbda != 1.0`, pointing at `GKDTrainer`, which
exposes the same knob for users who need on/off-policy mixing. The default
(`lmbda=1.0`, fully on-policy) does not warn.

Deprecate-before-remove: the removal lands in a later PR. Both policy modes are
already pinned end to end (see the on/off-policy training test), so that removal
will be a visible, guarded change.

The GKD (2306.13649) reproduction sets `lmbda=0.5` for the paper's on/off-policy
mixing. `DistillationTrainer` is dropping `lmbda` (becoming always on-policy),
so move the reproduction snippet to `GKDConfig` / `GKDTrainer`, which keep that
knob. Translate `max_completion_length` to GKD's `max_new_tokens`.

The prose now frames GKDTrainer as the home for the paper's `lmbda` mixing and
DistillationTrainer as the same generalized-JSD objective for the always-on-policy
case. Must land before `lmbda` is removed from DistillationConfig, so the index
never references a removed argument.

The trainer is now always on-policy: every generation-batch slice is generated
by the student, and there is no dataset-completion (off-policy) path.

Trainer:
- `_fill_buffer` generates for every slice; the per-slice on/off-policy coin flip
  (`random.random() <= lmbda`, broadcast across ranks) and the off-policy
  "store the dataset slice directly" branch are gone.
- Remove `self.lmbda`, `self._buffered_on_policy_flags`, and the four on/off-policy
  loss accumulators. `training_step` drops the on/off-policy attribution (its
  completion-length metrics lose the `on_policy_`/`off_policy_` prefix), and `log`
  drops the `on_policy_loss` / `off_policy_loss` emission and the all-reduce of
  the accumulator vector.
- Drop the now-unused `random` coin flip, `broadcast_object_list`, and
  `torch.distributed` / `DistributedType` imports (rich still uses `random`).

Config:
- Remove `lmbda`, its range check, the deprecation warning added last PR, and the
  `num_generations > 1 and lmbda < 1.0` warning.

Tests:
- `_make_args` drops `lmbda=0.0`; the on/off-policy training test becomes a single
  always-on-policy `test_train_updates_params`; the `lmbda` deprecation tests are
  removed; the server ragged-grad test drops `lmbda=0.0`.

Both policy modes were pinned before this (PR 02 / PR 10), so the off-policy
removal is a deliberate, test-visible deletion. Full experimental suite green.

…arse-top1

…arse-top1

The model card built by IWOPDTrainer attributed runs to the OPD paper
(2306.13649) carried over from the fork source. Point _paper at the
IW-OPD paper (2606.22600, "On the Position Bias of On-Policy
Distillation") so published checkpoints cite the correct work.

…arse-top1

…es-support

…es-support

…es-support

…es-support

…es-support

…pport

# Conflicts:
#	tests/experimental/test_distillation_trainer.py